from aiohttp import web, WSMsgType
import json
import logging
import asyncio
from datetime import datetime

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class WebSocketServer:
    def __init__(self):
        self.app = web.Application()
        self.clients = {}
        self.setup_routes()
        
    def setup_routes(self):
        self.app.router.add_get('/ws', self.websocket_handler)
        
    async def websocket_handler(self, request):
        ws = web.WebSocketResponse()
        await ws.prepare(request)
        
        client_id = str(id(ws))
        self.clients[client_id] = {
            'ws': ws,
            'connected_at': datetime.now().isoformat()
        }
        
        logger.info(f'Client {client_id} connected')
        
        try:
            async for msg in ws:
                if msg.type == WSMsgType.TEXT:
                    try:
                        data = json.loads(msg.data)
                        await self.handle_message(client_id, data)
                    except json.JSONDecodeError:
                        logger.error(f'Invalid JSON from client {client_id}')
                elif msg.type == WSMsgType.ERROR:
                    logger.error(f'WebSocket error: {ws.exception()}')
                    
        finally:
            await self.handle_disconnect(client_id)
            
        return ws
    
    async def handle_message(self, client_id, data):
        """处理接收到的消息"""
        msg_type = data.get('type')
        content = data.get('content')
        
        if msg_type == 'ping':
            await self.send_message(client_id, {
                'type': 'pong',
                'content': 'pong'
            })
        elif msg_type == 'broadcast':
            await self.broadcast_message(content, exclude_client=client_id)
        else:
            logger.info(f'Received message from {client_id}: {data}')
            
    async def send_message(self, client_id, message):
        """发送消息给指定客户端"""
        if client_id in self.clients:
            ws = self.clients[client_id]['ws']
            try:
                await ws.send_json(message)
            except Exception as e:
                logger.error(f'Error sending message to client {client_id}: {e}')
                
    async def broadcast_message(self, message, exclude_client=None):
        """广播消息给所有客户端"""
        for client_id in self.clients:
            if client_id != exclude_client:
                await self.send_message(client_id, {
                    'type': 'broadcast',
                    'content': message
                })
                
    async def handle_disconnect(self, client_id):
        """处理客户端断开连接"""
        if client_id in self.clients:
            logger.info(f'Client {client_id} disconnected')
            del self.clients[client_id]
            
    def run(self, host='0.0.0.0', port=8080):
        web.run_app(self.app, host=host, port=port)

if __name__ == '__main__':
    server = WebSocketServer()
    server.run() 